import os
from cubetools.python_video_frontend import PythonVideoFrontend


class PythonFrontend(object):

    # 在__init__中定义Python前端界面
    def __init__(self):
        self.demo = PythonVideoFrontend(
            model_name_cn='路面坑洼检测分割(yolov8s)',
            model_name_en='hf_yolov8s-pothole-segmentation',
            local_image=True,
            local_video=True,
            streaming_video=True,
            show_image_results=False,
            # results2text=lambda x: x[0],
            image_examples=['demo_data/p1.jpg', 'demo_data/p2.jpg', 'demo_data/p3.jpg', 'demo_data/p4.jpg', 'demo_data/p5.jpg'],
            local_video_examples=['demo_data/p1.mp4'],
            streaming_video_examples=[f'file://{os.getcwd()}/demo_data/p1.mp4'],
            readme='README.md'
        )

    def launch(self, **kwargs):
        self.demo.launch(**kwargs)
